import cplex
import random

class fs(cplex.Cplex):
    def __init__(self):
        cplex.Cplex.__init__(self) #has to be called to activa
        self.NI = 1; self.NS = 1
        self.I = range(self.NI); self.J = range(self.NS)
        self.x_name = []; self.y_name = []; self.rhs = []
        self.num_M = 0; self.M_rowname = []; self.M_colname = []; self.M_coef = []
        self.M_val = []; self.M_rank = []
        self.dm_lists = []; self.dm_diff_lists = []
        self.RLB = 0
        self.EPS = 0.00001
        self.times_called = 0
        self.flag = 1
        
    def gen_inst(self, NI=3, NS=3, RLB = 2, seed = 500):
        random.seed(seed)
        self.NI = NI;  self.NS = NS;
        self.RLB = RLB
        self.I = range(NI);  self.S = range(NS)
        self.x_name = ["x" + '_' + str(i) for i in self.I]
        self.y_name = ["y" + '_' + str(s) for s in self.S]
        #self.rhs = [[random.normalvariate(5,1) for i in self.I] for s in self.S]
        self.rhs = [[random.uniform(3,6) for i in self.I] for s in self.S]
        demand = self.rhs
        #cost = [random.uniform(0,1) for i in self.I]
        cost = [1]*NI
        
        x_name = self.x_name; x_type = ["C"] * NI
        x_lb   = [0] * NI; x_obj  = cost
        y_name = self.y_name; y_type = ["B"] * NS
        y_lb   = [0] * NS; y_obj  = [0] * NS

        # Add variables in one pass
        self.variables.add(
            obj = x_obj +  y_obj,
            types = x_type + y_type,
            names = x_name + y_name,
            lb = x_lb + y_lb)

        # Add knapsack constraints
        self.linear_constraints.add(
            lin_expr = [[y_name, [1]*NS]],
            senses = ["L"],
            rhs = [self.NS - self.RLB],  #this number means the number scenario to be dropped
            names = ["KP"])

        # capacity constraint
        self.dm_name = ["dm_" + str(i) + '_s_' +str(s) for s in self.S for i in self.I]
        self.linear_constraints.add(
            lin_expr = [ [ [x_name[i], y_name[s]], [1, demand[s][i]] ] for s in self.S for i in self.I],
            senses =   ["G"] * NI * NS,
            rhs = [ demand[s][i] for s in self.S for i in self.I],
            names = self.dm_name)

        # set M names
        self.num_M = self.NI
        self.M_ind = [[[0] for s in self.S] for i in self.I] 
        self.M_rowname = [ ["dm_" + str(i) + '_s_' + str(s) for s in self.S] for i in self.I ]
        self.M_colname = [ ['y_'+str(s) for s in self.S] for i in self.I]
        self.M_coef = self.linear_constraints.get_coefficients(zip(self.M_rowname[0], self.M_colname[0]))
        self.write("ccfs.lp")
        self.init_lists()


    def init_lists(self):
        self.dm_lists =      [[0.0 for s in self.S] for i in self.I]
        self.dm_diff_lists = [[0.0 for s in range(self.NS - 1)] for i in self.I]

        for i in self.I:
            self.dm_lists[i] = sorted(self.linear_constraints.get_rhs(self.M_rowname[i]), reverse=True)
            for s in range(self.NS -1 ):
                # the largest minus the following rank to get the difference
                self.dm_diff_lists[i][s] = self.dm_lists[i][0] - self.dm_lists[i][s+1]

    
    def set_bigm(self, rank = [1]):
        if len(rank) == 1:
            self.M_rank = [rank[0] for i in self.I] 
        elif len(rank) == self.NI:
            self.M_rank = list(rank)
        else:
            print "Wrong rank length"
            return
        
        self.M_val = [self.EPS + self.dm_diff_lists[i][self.M_rank[i]-1] for i in self.I]        

        M_val = self.M_val
        print M_val
        # Set bigm values
        for i in self.I:
            self.linear_constraints.set_coefficients(zip(self.M_rowname[i], self.M_colname[i], [M_val[i]]*self.NS ))
        print "Reset bigm values"
        self.write("ccfs_mod.lp")

    #This is the tricky part, we need to create overloaded
    #register_call function to call the register_call function
    # of the cplex version to pass data structures to them, 
    def register_callback(self, mycallback):        
        cb = cplex.Cplex.register_callback(self, mycallback)
        cb.fc = self




